#include "RAGSearch.h"


using namespace std;

RAGSearch::RAGSearch()
{
}

RAGSearch::~RAGSearch()
{
}

string RAGSearch::execPython(const std::string& cmd) {
    std::array<char, 256> buffer{};
    std::string result;

    // 打开管道，执行命令
    std::unique_ptr<FILE, decltype(&pclose)> pipe(popen(cmd.c_str(), "r"), pclose);
    if (!pipe) {
        throw std::runtime_error("popen() failed!");
    }

    // 读取输出
    while (fgets(buffer.data(), buffer.size(), pipe.get()) != nullptr) {
        result += buffer.data();
    }

    return result;
}


string RAGSearch::get_RAG_Handeled_Text(string user_query, vector<string> &all_searched_file_name) {

    embed.embed_str_data(user_query);

    std::string pythonCmd = "python3 ../utils/05_data_recall.py";
    std::string output;
    std::string enhanced_query;

    try {
        output = execPython(pythonCmd);

        // 处理检索到的内容
        enhanced_query = "我的问题是：\n" + user_query + "\n" + "我查找了相关的文档，内容如下，请你帮我整合一下并简要回答我的问题。\n";
        std::istringstream stream(output);
        std::string block;
        std::vector<std::string> doc_blocks;

        // 按 "---" 分割文档块
        std::string temp;
        while (std::getline(stream, temp)) {
            if (temp == "---") {
                if (!block.empty()) {
                    doc_blocks.push_back(block);
                    block.clear();
                }
            } else {
                if (!block.empty()) block += "\n";
                block += temp;
            }
        }
        if (!block.empty()) {
            doc_blocks.push_back(block);
        }

        std::unordered_set<std::string> file_name_set;


        // 处理每个文档块
        int part_num = 1;
        for (const auto& doc : doc_blocks) {
            std::istringstream blockStream(doc);
            std::string file_name;
            if (!std::getline(blockStream, file_name)) continue;

            // 去重
            if (file_name_set.find(file_name) == file_name_set.end()) {
                all_searched_file_name.push_back(file_name);
                file_name_set.insert(file_name);
            }

            enhanced_query += "第" + std::to_string(part_num) + "部分：" + file_name + "\n";
            std::string line;
            while (std::getline(blockStream, line)) {
                enhanced_query += line + "\n";
            }
            enhanced_query += "---\n";
            part_num++;
        }

        // 如果需要返回增强后的内容，返回 enhanced_query
        // return enhanced_query;

    } catch (const std::exception& ex) {
        std::cerr << "执行Python脚本失败: " << ex.what() << std::endl;
    }
    return enhanced_query;
}